import os
import pandas as pd
import numpy as np
from config import Config
from sklearn.model_selection import train_test_split


def load_data():
    """加载原始数据"""
    users = pd.read_csv(Config.USERS_FILE)
    products = pd.read_csv(Config.PRODUCTS_FILE)
    interactions = pd.read_csv(Config.INTERACTIONS_FILE)
    return users, products, interactions


def preprocess_data(users, products, interactions):
    """数据预处理"""
    # 转换行为为数值
    action_mapping = {'view': 1, 'like': 2, 'cart': 3, 'buy': 5}
    interactions['action_value'] = interactions['action'].map(action_mapping)

    # 添加时间权重（最近行为权重更高）
    max_timestamp = interactions['timestamp'].max()
    interactions['recency'] = (max_timestamp - interactions['timestamp']).dt.days
    interactions['time_weight'] = np.exp(-interactions['recency'] / 30)  # 30天衰减

    # 计算加权交互分数
    interactions['weighted_action'] = interactions['action_value'] * interactions['time_weight']

    # 合并数据
    data = interactions.merge(users, on='user_id')
    data = data.merge(products, on='product_id')

    # 划分训练测试集
    train, test = train_test_split(
        interactions[['user_id', 'product_id', 'weighted_action']],
        test_size=0.2,
        random_state=42
    )

    return data, train, test


def save_processed_data(data, train, test):
    """保存处理后的数据"""
    os.makedirs(Config.PROCESSED_DATA_DIR, exist_ok=True)

    data.to_pickle(os.path.join(Config.PROCESSED_DATA_DIR, "full_data.pkl"))
    train.to_csv(os.path.join(Config.PROCESSED_DATA_DIR, "train_interactions.csv"), index=False)
    test.to_csv(os.path.join(Config.PROCESSED_DATA_DIR, "test_interactions.csv"), index=False)

    print(f"Processed data saved to {Config.PROCESSED_DATA_DIR}")


def prepare_data():
    """数据准备流程"""
    users, products, interactions = load_data()

    # 转换时间格式
    interactions['timestamp'] = pd.to_datetime(interactions['timestamp'])

    data, train, test = preprocess_data(users, products, interactions)
    save_processed_data(data, train, test)

    return data, train, test